package com.sts.sink.impl;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.lang.Assert;
import cn.hutool.core.lang.Singleton;
import cn.hutool.core.util.StrUtil;
import com.sts.common.threadpool.VirtualThreadPool;
import com.sts.common.utils.JdbcConnectionUtil;
import com.sts.jdbc.JdbcOperation;
import com.sts.jdbc.model.ColumnInfo;
import com.sts.sink.StsSink;
import com.sts.sink.model.JdbcSinkConfig;
import com.sts.source.StsSource;
import com.sts.task.StsTask;
import com.sts.transform.StsTransform;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Semaphore;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Gatherers;
import java.util.stream.Stream;

@Slf4j
@Getter
public abstract class JdbcSinkImpl implements StsSink {

    private final StsSource stsSource;
    private final JdbcSinkConfig jdbcSinkConfig;
    private final StsTask stsTask;
    private final JdbcOperation jdbcOperation;

    public JdbcSinkImpl(StsSource stsSource, JdbcSinkConfig jdbcSinkConfig, StsTask stsTask, JdbcOperation jdbcOperation) {
        this.stsSource = stsSource;
        this.jdbcSinkConfig = jdbcSinkConfig;
        this.stsTask = stsTask;
        this.jdbcOperation = jdbcOperation;
    }

    // region insert
    @Override
    public void insert() {
        // 1. 获取来源头 b，c，a
        List<String> header = stsSource.getHeader();
        // 2. 获取 sink 目标表的所有字段信息
        Map<String, ColumnInfo> tableAllColumnMap = jdbcOperation.getTableAllColumn(
                JdbcConnectionUtil.getConnection(jdbcSinkConfig.getJdbcConnectionInfo()),
                jdbcSinkConfig.getSchema(),
                jdbcSinkConfig.getTable()
        );

        // 3. 校验头和表的字段信息是否匹配
        List<ColumnInfo> columnInfoList = checkHeaderMatchTableColumnAndReturnHeaderSortList(header, tableAllColumnMap);

        // 4. 构建 insert 语句
        String insertSql = buildInsertSql(columnInfoList);
        log.info("insert sql:{}", insertSql);

        // 5. jdbc入库
        int threadCount = jdbcSinkConfig.getThreadCount();
        if (threadCount <= 1) {
            log.info("单线程 insert");
            insertBySingleThread(columnInfoList, insertSql);
        } else {
            log.info("多线程 insert");
            insertByMultiThread(columnInfoList, insertSql);
        }
    }

    public void insertByMultiThread(List<ColumnInfo> columnInfoList, String insertSql) {
        try {
            int threadCount = jdbcSinkConfig.getThreadCount();
            Semaphore semaphore = new Semaphore(threadCount);
            List<CompletableFuture<Void>> futureArrayList = new ArrayList<>();

            Stream<List<String>> dataStream = stsSource.getDataStream();

            if (CollectionUtil.isNotEmpty(stsTask.getStsTransformList())) {
                for (StsTransform<?> stsTransform : stsTask.getStsTransformList()) {
                    dataStream = stsTransform.transform(dataStream);
                }
            }

            dataStream.takeWhile(_ -> stsTask.isFlag())
                    .gather(Gatherers.windowFixed(jdbcSinkConfig.getBatchSize()))
                    .forEach(dataList -> {
                        if (stsTask.isFlag()) {
                            try {
                                semaphore.acquire();
                            } catch (InterruptedException e) {
                                throw new RuntimeException(e.getMessage(), e);
                            }

                            CompletableFuture<Void> voidCompletableFuture = CompletableFuture.runAsync(
                                            () ->
                                                    insertDataSingleBatchSize(dataList, columnInfoList, insertSql),
                                            Singleton.get(VirtualThreadPool.class).getExecutorService()
                                    )
                                    .whenComplete((_, throwable) -> {
                                        if (throwable != null) {
                                            stsTask.setFlag(false);
                                        }
                                        semaphore.release();
                                    });
                            futureArrayList.add(voidCompletableFuture);
                        }
                    });

            if (CollectionUtil.isNotEmpty(futureArrayList)) {
                CompletableFuture.allOf(futureArrayList.toArray(new CompletableFuture[0])).join();
            }
        } finally {
            try {
                stsSource.close();
            } catch (Exception e) {
                log.error("关闭 source 失败: {}", e.getMessage(), e);
            }
        }
    }

    public void insertDataSingleBatchSize(List<List<String>> dataList, List<ColumnInfo> columnInfoList, String insertSql) {
        try (
                Connection connection = JdbcConnectionUtil.getConnection(jdbcSinkConfig.getJdbcConnectionInfo());
                PreparedStatement preparedStatement = connection.prepareStatement(insertSql)
        ) {
            for (List<String> rowValueList : dataList) {
                if (!stsTask.isFlag()) {
                    break;
                }
                insertSetCellValue(preparedStatement, columnInfoList, rowValueList);
                preparedStatement.addBatch();
            }
            int[] successArr = preparedStatement.executeBatch();
            stsTask.updateTotalSuccessCount(stsTask, successArr.length);
            preparedStatement.clearBatch();
        } catch (Exception e) {
            throw new RuntimeException(e.getMessage(), e);
        }
    }

    public void insertBySingleThread(List<ColumnInfo> columnInfoList, String insertSql) {
        try {
            try (
                    Connection connection = JdbcConnectionUtil.getConnection(jdbcSinkConfig.getJdbcConnectionInfo());
                    PreparedStatement preparedStatement = connection.prepareStatement(insertSql)
            ) {
                Stream<List<String>> dataStream = stsSource.getDataStream();

                if (CollectionUtil.isNotEmpty(stsTask.getStsTransformList())) {
                    for (StsTransform<?> stsTransform : stsTask.getStsTransformList()) {
                        dataStream = stsTransform.transform(dataStream);
                    }
                }

                dataStream.takeWhile(_ -> stsTask.isFlag())
                        .gather(Gatherers.windowFixed(jdbcSinkConfig.getBatchSize()))
                        .forEach(dataList -> {
                            try {
                                for (List<String> rowValue : dataList) {
                                    insertSetCellValue(preparedStatement, columnInfoList, rowValue);
                                    preparedStatement.addBatch();
                                }

                                int[] successArr = preparedStatement.executeBatch();
                                stsTask.updateTotalSuccessCount(stsTask, successArr.length);
                                preparedStatement.clearBatch();
                            } catch (Exception e) {
                                throw new RuntimeException(e.getMessage(), e);
                            }
                        });
            } catch (SQLException e) {
                throw new RuntimeException(e.getMessage(), e);
            }
        } finally {
            try {
                stsSource.close();
            } catch (Exception e) {
                log.error("关闭 source 失败: {}", e.getMessage(), e);
            }
        }
    }

    public abstract void insertSetCellValue(PreparedStatement preparedStatement, List<ColumnInfo> columnInfoList, List<String> rowValueList) throws Exception;

    public String buildInsertSql(List<ColumnInfo> columnInfoList) {
        String identifier = jdbcSinkConfig.getJdbcConnectionInfo().getIdentifier();

        StringBuilder insertSql = new StringBuilder("INSERT INTO ");
        insertSql.append(identifier).append(jdbcSinkConfig.getSchema()).append(identifier)
                .append(".")
                .append(identifier).append(jdbcSinkConfig.getTable()).append(identifier)
                .append("(");

        for (ColumnInfo columnInfo : columnInfoList) {
            insertSql.append(identifier).append(columnInfo.getColumnName()).append(identifier).append(", ");
        }

        insertSql.delete(insertSql.lastIndexOf(","), insertSql.length());
        insertSql.append(") VALUES (");

        insertSql.append("?, ".repeat(columnInfoList.size()));

        insertSql.delete(insertSql.lastIndexOf(","), insertSql.length());
        insertSql.append(")");

        return insertSql.toString();
    }

    public List<ColumnInfo> checkHeaderMatchTableColumnAndReturnHeaderSortList(List<String> headerList, Map<String, ColumnInfo> tableAllColumnMap) {
        // 检验来源头中是否存在重复的列
        StringBuilder repeatHeaderErrorMsg = new StringBuilder();
        // 检验来源和目标字段差异
        StringBuilder matchColumnErrorMsg = new StringBuilder();

        Set<String> set = new HashSet<>();
        List<ColumnInfo> resultList = new ArrayList<>();

        for (String colName : headerList) {
            if (!set.add(colName)) {
                repeatHeaderErrorMsg.append(colName).append(", ");
            }

            String colNameUpperCase = colName.toUpperCase();
            if (!tableAllColumnMap.containsKey(colNameUpperCase)) {
                matchColumnErrorMsg.append(colName).append(", ");
            } else {
                resultList.add(tableAllColumnMap.get(colNameUpperCase));
            }
        }

        if (StrUtil.isNotBlank(repeatHeaderErrorMsg)) {
            String repeatHeaderErrorMsgStr = StrUtil.format("来源中存在重复列: {}",
                    repeatHeaderErrorMsg.delete(repeatHeaderErrorMsg.lastIndexOf(","), repeatHeaderErrorMsg.length()));
            throw new RuntimeException(repeatHeaderErrorMsgStr);
        }

        if (StrUtil.isNotBlank(matchColumnErrorMsg)) {
            String matchColumnErrorMsgStr = StrUtil.format("{} 字段在表中不存在",
                    matchColumnErrorMsg.delete(matchColumnErrorMsg.lastIndexOf(","), matchColumnErrorMsg.length()));
            throw new RuntimeException(matchColumnErrorMsgStr);
        }

        return resultList;
    }
    // endregion

    // region upsert
    @Override
    public void upsert() {
        // 1. 获取头信息
        List<String> header = stsSource.getHeader();

        // 2. 获取 sink 目标表的所有字段信息
        Map<String, ColumnInfo> tableAllColumnMap = jdbcOperation.getTableAllColumn(
                JdbcConnectionUtil.getConnection(jdbcSinkConfig.getJdbcConnectionInfo()),
                jdbcSinkConfig.getSchema(),
                jdbcSinkConfig.getTable()
        );

        // 3. 校验头信息和目标表的字段信息是否有出入,
        // 并返回数据库中的列的字段信息
        // (防止数据库区分列字段的大小写, 因此直接使用字段信息中的列名)
        List<ColumnInfo> columnInfoList = checkHeaderMatchTableColumnAndReturnHeaderSortList(header, tableAllColumnMap);

        // 4. 获取需要导入的表主键字段信息
        Map<String, ColumnInfo> tableAllPrimaryKeyColumnMap = jdbcOperation.getTableAllPrimaryKeyColumn(
                JdbcConnectionUtil.getConnection(jdbcSinkConfig.getJdbcConnectionInfo()),
                jdbcSinkConfig.getSchema(),
                jdbcSinkConfig.getTable()
        );
        Assert.isTrue(CollectionUtil.isNotEmpty(tableAllPrimaryKeyColumnMap), "未检测到表中的主键, Upsert需要主键的支持~");

        // 5. 校验主键是否在source header中存在
        checkSourcePrimaryKeyColumn(columnInfoList, tableAllPrimaryKeyColumnMap);

        // 6. 构建upsert sql
        String upsertSql = buildUpsertSql(columnInfoList, tableAllPrimaryKeyColumnMap);
        log.info("upsertSql: {}", upsertSql);

        // 7. jdbc入库
        int threadCount = jdbcSinkConfig.getThreadCount();
        if (threadCount <= 1) {
            log.info("单线程 upsert");
            upsertBySingleThread(columnInfoList, upsertSql, tableAllPrimaryKeyColumnMap);
        } else {
            log.info("多线程 upsert");
            upsertByMultiThread(columnInfoList, upsertSql, tableAllPrimaryKeyColumnMap);
        }
    }

    public void upsertByMultiThread(List<ColumnInfo> columnInfoList, String upsertSql, Map<String, ColumnInfo> tableAllPrimaryKeyColumnMap) {
        try {
            int threadCount = jdbcSinkConfig.getThreadCount();
            Semaphore semaphore = new Semaphore(threadCount);
            List<CompletableFuture<Void>> futureArrayList = new ArrayList<>();

            Stream<List<String>> dataStream = stsSource.getDataStream();

            if (CollectionUtil.isNotEmpty(stsTask.getStsTransformList())) {
                for (StsTransform<?> stsTransform : stsTask.getStsTransformList()) {
                    dataStream = stsTransform.transform(dataStream);
                }
            }

            dataStream.takeWhile(_ -> stsTask.isFlag())
                    .gather(Gatherers.windowFixed(jdbcSinkConfig.getBatchSize()))
                    .forEach(dataList -> {
                        if (stsTask.isFlag()) {
                            try {
                                semaphore.acquire();
                            } catch (InterruptedException e) {
                                throw new RuntimeException(e.getMessage(), e);
                            }

                            CompletableFuture<Void> voidCompletableFuture = CompletableFuture.runAsync(
                                            () ->
                                                    upsertDataSingleBatchSize(dataList, columnInfoList, upsertSql, tableAllPrimaryKeyColumnMap),
                                            Singleton.get(VirtualThreadPool.class).getExecutorService()
                                    )
                                    .whenComplete((_, throwable) -> {
                                        if (throwable != null) {
                                            stsTask.setFlag(false);
                                        }
                                        semaphore.release();
                                    });
                            futureArrayList.add(voidCompletableFuture);
                        }
                    });

            if (CollectionUtil.isNotEmpty(futureArrayList)) {
                CompletableFuture.allOf(futureArrayList.toArray(new CompletableFuture[0])).join();
            }
        } finally {
            try {
                stsSource.close();
            } catch (Exception e) {
                log.error("关闭 source 失败: {}", e.getMessage(), e);
            }
        }
    }

    public void upsertDataSingleBatchSize(List<List<String>> dataList, List<ColumnInfo> columnInfoList, String upsertSql, Map<String, ColumnInfo> tableAllPrimaryKeyColumnMap) {
        try (
                Connection connection = JdbcConnectionUtil.getConnection(jdbcSinkConfig.getJdbcConnectionInfo());
                PreparedStatement preparedStatement = connection.prepareStatement(upsertSql)
        ) {
            for (List<String> rowValueList : dataList) {
                if (!stsTask.isFlag()) {
                    break;
                }
                upsertSetCellValue(preparedStatement, columnInfoList, rowValueList, tableAllPrimaryKeyColumnMap);
                preparedStatement.addBatch();
            }
            int[] successArr = preparedStatement.executeBatch();
            stsTask.updateTotalSuccessCount(stsTask, successArr.length);
            preparedStatement.clearBatch();
        } catch (Exception e) {
            throw new RuntimeException(e.getMessage(), e);
        }
    }

    public void upsertBySingleThread(List<ColumnInfo> columnInfoList, String upsertSql, Map<String, ColumnInfo> tableAllPrimaryKeyColumnMap) {
        try {
            try (
                    Connection connection = JdbcConnectionUtil.getConnection(jdbcSinkConfig.getJdbcConnectionInfo());
                    PreparedStatement preparedStatement = connection.prepareStatement(upsertSql)
            ) {
                Stream<List<String>> dataStream = stsSource.getDataStream();

                if (CollectionUtil.isNotEmpty(stsTask.getStsTransformList())) {
                    for (StsTransform<?> stsTransform : stsTask.getStsTransformList()) {
                        dataStream = stsTransform.transform(dataStream);
                    }
                }

                dataStream.takeWhile(_ -> stsTask.isFlag())
                        .gather(Gatherers.windowFixed(jdbcSinkConfig.getBatchSize()))
                        .forEach(dataList -> {
                            try {
                                for (List<String> rowValue : dataList) {
                                    upsertSetCellValue(preparedStatement, columnInfoList, rowValue, tableAllPrimaryKeyColumnMap);
                                    preparedStatement.addBatch();
                                }

                                int[] successArr = preparedStatement.executeBatch();
                                stsTask.updateTotalSuccessCount(stsTask, successArr.length);
                                preparedStatement.clearBatch();
                            } catch (Exception e) {
                                throw new RuntimeException(e.getMessage(), e);
                            }
                        });
            } catch (SQLException e) {
                throw new RuntimeException(e.getMessage(), e);
            }
        } finally {
            try {
                stsSource.close();
            } catch (Exception e) {
                log.error("关闭 source 失败: {}", e.getMessage(), e);
            }
        }
    }

    public void checkSourcePrimaryKeyColumn(List<ColumnInfo> columnInfoList, Map<String, ColumnInfo> tableAllPrimaryKeyColumnMap) {
        StringBuilder stringBuilder = new StringBuilder();

        // key: 字段名大写
        Map<String, ColumnInfo> columnInfoModelMap = columnInfoList.stream()
                .collect(Collectors.toMap(
                        columnInfo -> columnInfo.getColumnName().toUpperCase(),
                        Function.identity())
                );

        // 遍历主键map
        for (Map.Entry<String, ColumnInfo> entry : tableAllPrimaryKeyColumnMap.entrySet()) {
            String key = entry.getKey();
            ColumnInfo value = entry.getValue();
            // 如果 source 中不存在说明有问题, 需要记录后抛异常
            if (!columnInfoModelMap.containsKey(key)) {
                stringBuilder.append(value.getColumnName()).append(", ");
            }
        }

        if (StrUtil.isNotBlank(stringBuilder)) {
            stringBuilder.delete(stringBuilder.lastIndexOf(","), stringBuilder.length());
            throw new RuntimeException("Upsert 需要主键字段都在 source 存在, 以下列在 source 中不存在: " + stringBuilder);
        }
    }

    public abstract String buildUpsertSql(List<ColumnInfo> columnInfoList, Map<String, ColumnInfo> tableAllPrimaryKeyColumnMap);

    public abstract void upsertSetCellValue(PreparedStatement preparedStatement, List<ColumnInfo> columnInfoList, List<String> rowValueList, Map<String, ColumnInfo> tableAllPrimaryKeyColumnMap) throws Exception;
    // endregion
}
